{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "f1dfd015",
      "metadata": {
        "id": "f1dfd015"
      },
      "source": [
        "# Drug Recommendation using MICRON Model on MIMIC-III Dataset\n",
        "\n",
        "This notebook demonstrates how to use the MICRON model for drug recommendation using the MIMIC-III dataset. The model is implemented using PyHealth 2.0 framework.\n",
        "\n",
        "MICRON (Medication reCommendation using Recurrent ResIdual Networks) is designed to predict medications based on patient diagnoses and procedures."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "46345c0b",
      "metadata": {
        "id": "46345c0b"
      },
      "source": [
        "## 1. Setup Google Drive and Environment\n",
        "\n",
        "First, we'll mount Google Drive to access and save our data. We'll also install PyHealth from the forked repository and its dependencies. The notebook uses the latest version of PyHealth from https://github.com/naveenkcb/PyHealth."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "1f130a36",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "1f130a36",
        "outputId": "f4f6b214-6dfa-49ca-f0cd-98e7ccf1905a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting git+https://github.com/naveenkcb/PyHealth.git\n",
            "  Cloning https://github.com/naveenkcb/PyHealth.git to /tmp/pip-req-build-e8mrcm25\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/naveenkcb/PyHealth.git /tmp/pip-req-build-e8mrcm25\n",
            "  Resolved https://github.com/naveenkcb/PyHealth.git to commit bb3da26de8c9747fd67b6842096a9a07b60310eb\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.12/dist-packages (from pyhealth==2.0a8) (1.11.0)\n",
            "Collecting mne~=1.10.0 (from pyhealth==2.0a8)\n",
            "  Downloading mne-1.10.2-py3-none-any.whl.metadata (21 kB)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.12/dist-packages (from pyhealth==2.0a8) (3.5)\n",
            "Collecting numpy~=1.26.4 (from pyhealth==2.0a8)\n",
            "  Downloading numpy-1.26.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (61 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m61.0/61.0 kB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting pandarallel~=1.6.5 (from pyhealth==2.0a8)\n",
            "  Downloading pandarallel-1.6.5.tar.gz (14 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting pandas~=2.3.1 (from pyhealth==2.0a8)\n",
            "  Downloading pandas-2.3.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.metadata (91 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m91.2/91.2 kB\u001b[0m \u001b[31m9.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: peft in /usr/local/lib/python3.12/dist-packages (from pyhealth==2.0a8) (0.17.1)\n",
            "Collecting polars~=1.31.0 (from pyhealth==2.0a8)\n",
            "  Downloading polars-1.31.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (14 kB)\n",
            "Requirement already satisfied: pydantic~=2.11.7 in /usr/local/lib/python3.12/dist-packages (from pyhealth==2.0a8) (2.11.10)\n",
            "Collecting rdkit (from pyhealth==2.0a8)\n",
            "  Downloading rdkit-2025.9.1-cp312-cp312-manylinux_2_28_x86_64.whl.metadata (4.1 kB)\n",
            "Collecting scikit-learn~=1.7.0 (from pyhealth==2.0a8)\n",
            "  Downloading scikit_learn-1.7.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (11 kB)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.12/dist-packages (from pyhealth==2.0a8) (0.23.0+cu126)\n",
            "Collecting torch~=2.7.1 (from pyhealth==2.0a8)\n",
            "  Downloading torch-2.7.1-cp312-cp312-manylinux_2_28_x86_64.whl.metadata (29 kB)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.12/dist-packages (from pyhealth==2.0a8) (4.67.1)\n",
            "Collecting transformers~=4.53.2 (from pyhealth==2.0a8)\n",
            "  Downloading transformers-4.53.3-py3-none-any.whl.metadata (40 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m40.9/40.9 kB\u001b[0m \u001b[31m3.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: urllib3~=2.5.0 in /usr/local/lib/python3.12/dist-packages (from pyhealth==2.0a8) (2.5.0)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.12/dist-packages (from mne~=1.10.0->pyhealth==2.0a8) (4.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from mne~=1.10.0->pyhealth==2.0a8) (3.1.6)\n",
            "Requirement already satisfied: lazy-loader>=0.3 in /usr/local/lib/python3.12/dist-packages (from mne~=1.10.0->pyhealth==2.0a8) (0.4)\n",
            "Requirement already satisfied: matplotlib>=3.7 in /usr/local/lib/python3.12/dist-packages (from mne~=1.10.0->pyhealth==2.0a8) (3.10.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.12/dist-packages (from mne~=1.10.0->pyhealth==2.0a8) (25.0)\n",
            "Requirement already satisfied: pooch>=1.5 in /usr/local/lib/python3.12/dist-packages (from mne~=1.10.0->pyhealth==2.0a8) (1.8.2)\n",
            "Requirement already satisfied: scipy>=1.11 in /usr/local/lib/python3.12/dist-packages (from mne~=1.10.0->pyhealth==2.0a8) (1.16.3)\n",
            "Requirement already satisfied: dill>=0.3.1 in /usr/local/lib/python3.12/dist-packages (from pandarallel~=1.6.5->pyhealth==2.0a8) (0.3.8)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.12/dist-packages (from pandarallel~=1.6.5->pyhealth==2.0a8) (5.9.5)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas~=2.3.1->pyhealth==2.0a8) (2.9.0.post0)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas~=2.3.1->pyhealth==2.0a8) (2025.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas~=2.3.1->pyhealth==2.0a8) (2025.2)\n",
            "Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.12/dist-packages (from pydantic~=2.11.7->pyhealth==2.0a8) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.33.2 in /usr/local/lib/python3.12/dist-packages (from pydantic~=2.11.7->pyhealth==2.0a8) (2.33.2)\n",
            "Requirement already satisfied: typing-extensions>=4.12.2 in /usr/local/lib/python3.12/dist-packages (from pydantic~=2.11.7->pyhealth==2.0a8) (4.15.0)\n",
            "Requirement already satisfied: typing-inspection>=0.4.0 in /usr/local/lib/python3.12/dist-packages (from pydantic~=2.11.7->pyhealth==2.0a8) (0.4.2)\n",
            "Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn~=1.7.0->pyhealth==2.0a8) (1.5.2)\n",
            "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn~=1.7.0->pyhealth==2.0a8) (3.6.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (3.20.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (75.2.0)\n",
            "Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (1.13.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (2025.3.0)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (12.6.77)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (12.6.77)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.6.80 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (12.6.80)\n",
            "Collecting nvidia-cudnn-cu12==9.5.1.17 (from torch~=2.7.1->pyhealth==2.0a8)\n",
            "  Downloading nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl.metadata (1.6 kB)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.6.4.1 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (12.6.4.1)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.3.0.4 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (11.3.0.4)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.7.77 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (10.3.7.77)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.7.1.2 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (11.7.1.2)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.5.4.2 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (12.5.4.2)\n",
            "Collecting nvidia-cusparselt-cu12==0.6.3 (from torch~=2.7.1->pyhealth==2.0a8)\n",
            "  Downloading nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl.metadata (6.8 kB)\n",
            "Collecting nvidia-nccl-cu12==2.26.2 (from torch~=2.7.1->pyhealth==2.0a8)\n",
            "  Downloading nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (2.0 kB)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (12.6.77)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12==12.6.85 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (12.6.85)\n",
            "Requirement already satisfied: nvidia-cufile-cu12==1.11.1.6 in /usr/local/lib/python3.12/dist-packages (from torch~=2.7.1->pyhealth==2.0a8) (1.11.1.6)\n",
            "Collecting triton==3.3.1 (from torch~=2.7.1->pyhealth==2.0a8)\n",
            "  Downloading triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (1.5 kB)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.30.0 in /usr/local/lib/python3.12/dist-packages (from transformers~=4.53.2->pyhealth==2.0a8) (0.36.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.12/dist-packages (from transformers~=4.53.2->pyhealth==2.0a8) (6.0.3)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.12/dist-packages (from transformers~=4.53.2->pyhealth==2.0a8) (2024.11.6)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.12/dist-packages (from transformers~=4.53.2->pyhealth==2.0a8) (2.32.4)\n",
            "Collecting tokenizers<0.22,>=0.21 (from transformers~=4.53.2->pyhealth==2.0a8)\n",
            "  Downloading tokenizers-0.21.4-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.7 kB)\n",
            "Requirement already satisfied: safetensors>=0.4.3 in /usr/local/lib/python3.12/dist-packages (from transformers~=4.53.2->pyhealth==2.0a8) (0.6.2)\n",
            "Requirement already satisfied: Pillow in /usr/local/lib/python3.12/dist-packages (from rdkit->pyhealth==2.0a8) (11.3.0)\n",
            "INFO: pip is looking at multiple versions of torchvision to determine which version is compatible with other requirements. This could take a while.\n",
            "Collecting torchvision (from pyhealth==2.0a8)\n",
            "  Downloading torchvision-0.24.0-cp312-cp312-manylinux_2_28_x86_64.whl.metadata (5.9 kB)\n",
            "  Downloading torchvision-0.23.0-cp312-cp312-manylinux_2_28_x86_64.whl.metadata (6.1 kB)\n",
            "  Downloading torchvision-0.22.1-cp312-cp312-manylinux_2_28_x86_64.whl.metadata (6.1 kB)\n",
            "Requirement already satisfied: hf-xet<2.0.0,>=1.1.3 in /usr/local/lib/python3.12/dist-packages (from huggingface-hub<1.0,>=0.30.0->transformers~=4.53.2->pyhealth==2.0a8) (1.2.0)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib>=3.7->mne~=1.10.0->pyhealth==2.0a8) (1.3.3)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.12/dist-packages (from matplotlib>=3.7->mne~=1.10.0->pyhealth==2.0a8) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.12/dist-packages (from matplotlib>=3.7->mne~=1.10.0->pyhealth==2.0a8) (4.60.1)\n",
            "Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib>=3.7->mne~=1.10.0->pyhealth==2.0a8) (1.4.9)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib>=3.7->mne~=1.10.0->pyhealth==2.0a8) (3.2.5)\n",
            "Requirement already satisfied: platformdirs>=2.5.0 in /usr/local/lib/python3.12/dist-packages (from pooch>=1.5->mne~=1.10.0->pyhealth==2.0a8) (4.5.0)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas~=2.3.1->pyhealth==2.0a8) (1.17.0)\n",
            "Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests->transformers~=4.53.2->pyhealth==2.0a8) (3.4.4)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests->transformers~=4.53.2->pyhealth==2.0a8) (3.11)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests->transformers~=4.53.2->pyhealth==2.0a8) (2025.10.5)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch~=2.7.1->pyhealth==2.0a8) (1.3.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->mne~=1.10.0->pyhealth==2.0a8) (3.0.3)\n",
            "Downloading mne-1.10.2-py3-none-any.whl (7.4 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.4/7.4 MB\u001b[0m \u001b[31m76.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading numpy-1.26.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (18.0 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m18.0/18.0 MB\u001b[0m \u001b[31m66.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading pandas-2.3.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (12.4 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m12.4/12.4 MB\u001b[0m \u001b[31m143.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading polars-1.31.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (35.1 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m35.1/35.1 MB\u001b[0m \u001b[31m13.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading scikit_learn-1.7.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (9.5 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m9.5/9.5 MB\u001b[0m \u001b[31m91.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading torch-2.7.1-cp312-cp312-manylinux_2_28_x86_64.whl (821.0 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m821.0/821.0 MB\u001b[0m \u001b[31m2.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl (571.0 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m571.0/571.0 MB\u001b[0m \u001b[31m3.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl (156.8 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m156.8/156.8 MB\u001b[0m \u001b[31m9.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (201.3 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m201.3/201.3 MB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (155.7 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m155.7/155.7 MB\u001b[0m \u001b[31m7.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading transformers-4.53.3-py3-none-any.whl (10.8 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.8/10.8 MB\u001b[0m \u001b[31m83.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading rdkit-2025.9.1-cp312-cp312-manylinux_2_28_x86_64.whl (36.2 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m36.2/36.2 MB\u001b[0m \u001b[31m18.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading torchvision-0.22.1-cp312-cp312-manylinux_2_28_x86_64.whl (7.5 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.5/7.5 MB\u001b[0m \u001b[31m77.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading tokenizers-0.21.4-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.1 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m47.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hBuilding wheels for collected packages: pyhealth, pandarallel\n",
            "  Building wheel for pyhealth (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyhealth: filename=pyhealth-2.0a8-py3-none-any.whl size=369656 sha256=d9ab93303c6f94cd947522e31488863518d4ac8213781c3cfe3f101e4a716124\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-q8ibyuty/wheels/e9/10/11/3146f609c6b24edf823d697c4a93da2e447bada2d1fb3fb819\n",
            "  Building wheel for pandarallel (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pandarallel: filename=pandarallel-1.6.5-py3-none-any.whl size=16674 sha256=401f3b297277ed66c5c78a2f56e7aae941b2b53989461d6f90999d32004cb5af\n",
            "  Stored in directory: /root/.cache/pip/wheels/46/f9/0d/40c9cd74a7cb8dc8fe57e8d6c3c19e2c730449c0d3f2bf66b5\n",
            "Successfully built pyhealth pandarallel\n",
            "Installing collected packages: nvidia-cusparselt-cu12, triton, polars, nvidia-nccl-cu12, nvidia-cudnn-cu12, numpy, rdkit, pandas, torch, tokenizers, scikit-learn, pandarallel, transformers, torchvision, mne, pyhealth\n",
            "  Attempting uninstall: nvidia-cusparselt-cu12\n",
            "    Found existing installation: nvidia-cusparselt-cu12 0.7.1\n",
            "    Uninstalling nvidia-cusparselt-cu12-0.7.1:\n",
            "      Successfully uninstalled nvidia-cusparselt-cu12-0.7.1\n",
            "  Attempting uninstall: triton\n",
            "    Found existing installation: triton 3.4.0\n",
            "    Uninstalling triton-3.4.0:\n",
            "      Successfully uninstalled triton-3.4.0\n",
            "  Attempting uninstall: polars\n",
            "    Found existing installation: polars 1.25.2\n",
            "    Uninstalling polars-1.25.2:\n",
            "      Successfully uninstalled polars-1.25.2\n",
            "  Attempting uninstall: nvidia-nccl-cu12\n",
            "    Found existing installation: nvidia-nccl-cu12 2.27.3\n",
            "    Uninstalling nvidia-nccl-cu12-2.27.3:\n",
            "      Successfully uninstalled nvidia-nccl-cu12-2.27.3\n",
            "  Attempting uninstall: nvidia-cudnn-cu12\n",
            "    Found existing installation: nvidia-cudnn-cu12 9.10.2.21\n",
            "    Uninstalling nvidia-cudnn-cu12-9.10.2.21:\n",
            "      Successfully uninstalled nvidia-cudnn-cu12-9.10.2.21\n",
            "  Attempting uninstall: numpy\n",
            "    Found existing installation: numpy 2.0.2\n",
            "    Uninstalling numpy-2.0.2:\n",
            "      Successfully uninstalled numpy-2.0.2\n",
            "  Attempting uninstall: pandas\n",
            "    Found existing installation: pandas 2.2.2\n",
            "    Uninstalling pandas-2.2.2:\n",
            "      Successfully uninstalled pandas-2.2.2\n",
            "  Attempting uninstall: torch\n",
            "    Found existing installation: torch 2.8.0+cu126\n",
            "    Uninstalling torch-2.8.0+cu126:\n",
            "      Successfully uninstalled torch-2.8.0+cu126\n",
            "  Attempting uninstall: tokenizers\n",
            "    Found existing installation: tokenizers 0.22.1\n",
            "    Uninstalling tokenizers-0.22.1:\n",
            "      Successfully uninstalled tokenizers-0.22.1\n",
            "  Attempting uninstall: scikit-learn\n",
            "    Found existing installation: scikit-learn 1.6.1\n",
            "    Uninstalling scikit-learn-1.6.1:\n",
            "      Successfully uninstalled scikit-learn-1.6.1\n",
            "  Attempting uninstall: transformers\n",
            "    Found existing installation: transformers 4.57.1\n",
            "    Uninstalling transformers-4.57.1:\n",
            "      Successfully uninstalled transformers-4.57.1\n",
            "  Attempting uninstall: torchvision\n",
            "    Found existing installation: torchvision 0.23.0+cu126\n",
            "    Uninstalling torchvision-0.23.0+cu126:\n",
            "      Successfully uninstalled torchvision-0.23.0+cu126\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "google-colab 1.0.0 requires pandas==2.2.2, but you have pandas 2.3.3 which is incompatible.\n",
            "torchaudio 2.8.0+cu126 requires torch==2.8.0, but you have torch 2.7.1 which is incompatible.\n",
            "opencv-python-headless 4.12.0.88 requires numpy<2.3.0,>=2; python_version >= \"3.9\", but you have numpy 1.26.4 which is incompatible.\n",
            "cudf-polars-cu12 25.6.0 requires polars<1.29,>=1.25, but you have polars 1.31.0 which is incompatible.\n",
            "opencv-python 4.12.0.88 requires numpy<2.3.0,>=2; python_version >= \"3.9\", but you have numpy 1.26.4 which is incompatible.\n",
            "jax 0.7.2 requires numpy>=2.0, but you have numpy 1.26.4 which is incompatible.\n",
            "pytensor 2.35.1 requires numpy>=2.0, but you have numpy 1.26.4 which is incompatible.\n",
            "thinc 8.3.6 requires numpy<3.0.0,>=2.0.0, but you have numpy 1.26.4 which is incompatible.\n",
            "opencv-contrib-python 4.12.0.88 requires numpy<2.3.0,>=2; python_version >= \"3.9\", but you have numpy 1.26.4 which is incompatible.\n",
            "cudf-cu12 25.6.0 requires pandas<2.2.4dev0,>=2.0, but you have pandas 2.3.3 which is incompatible.\n",
            "jaxlib 0.7.2 requires numpy>=2.0, but you have numpy 1.26.4 which is incompatible.\n",
            "dask-cudf-cu12 25.6.0 requires pandas<2.2.4dev0,>=2.0, but you have pandas 2.3.3 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed mne-1.10.2 numpy-1.26.4 nvidia-cudnn-cu12-9.5.1.17 nvidia-cusparselt-cu12-0.6.3 nvidia-nccl-cu12-2.26.2 pandarallel-1.6.5 pandas-2.3.3 polars-1.31.0 pyhealth-2.0a8 rdkit-2025.9.1 scikit-learn-1.7.2 tokenizers-0.21.4 torch-2.7.1 torchvision-0.22.1 transformers-4.53.3 triton-3.3.1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "numpy"
                ]
              },
              "id": "95d99e6d90f043b7be6177524843f7df"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: torch in /usr/local/lib/python3.12/dist-packages (2.7.1)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.12/dist-packages (1.7.2)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.12/dist-packages (2.3.3)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.12/dist-packages (1.26.4)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.12/dist-packages (4.67.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from torch) (3.20.0)\n",
            "Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.12/dist-packages (from torch) (4.15.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from torch) (75.2.0)\n",
            "Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch) (1.13.3)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.12/dist-packages (from torch) (3.5)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from torch) (3.1.6)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.12/dist-packages (from torch) (2025.3.0)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch) (12.6.77)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch) (12.6.77)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.6.80 in /usr/local/lib/python3.12/dist-packages (from torch) (12.6.80)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.5.1.17 in /usr/local/lib/python3.12/dist-packages (from torch) (9.5.1.17)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.6.4.1 in /usr/local/lib/python3.12/dist-packages (from torch) (12.6.4.1)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.3.0.4 in /usr/local/lib/python3.12/dist-packages (from torch) (11.3.0.4)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.7.77 in /usr/local/lib/python3.12/dist-packages (from torch) (10.3.7.77)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.7.1.2 in /usr/local/lib/python3.12/dist-packages (from torch) (11.7.1.2)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.5.4.2 in /usr/local/lib/python3.12/dist-packages (from torch) (12.5.4.2)\n",
            "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.3 in /usr/local/lib/python3.12/dist-packages (from torch) (0.6.3)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.26.2 in /usr/local/lib/python3.12/dist-packages (from torch) (2.26.2)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.6.77 in /usr/local/lib/python3.12/dist-packages (from torch) (12.6.77)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12==12.6.85 in /usr/local/lib/python3.12/dist-packages (from torch) (12.6.85)\n",
            "Requirement already satisfied: nvidia-cufile-cu12==1.11.1.6 in /usr/local/lib/python3.12/dist-packages (from torch) (1.11.1.6)\n",
            "Requirement already satisfied: triton==3.3.1 in /usr/local/lib/python3.12/dist-packages (from torch) (3.3.1)\n",
            "Requirement already satisfied: scipy>=1.8.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn) (1.16.3)\n",
            "Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn) (1.5.2)\n",
            "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn) (3.6.0)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas) (2.9.0.post0)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas) (2025.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas) (2025.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas) (1.17.0)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch) (1.3.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->torch) (3.0.3)\n"
          ]
        }
      ],
      "source": [
        "# Mount Google Drive\n",
        "#from google.colab import drive\n",
        "#drive.mount('/content/drive')\n",
        "\n",
        "# Install PyHealth from your forked repository\n",
        "!pip install git+https://github.com/naveenkcb/PyHealth.git\n",
        "# Install other required packages\n",
        "!pip install torch scikit-learn pandas numpy tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f70c5176",
      "metadata": {
        "id": "f70c5176"
      },
      "source": [
        "## 2. Import Required Libraries and Setup Configuration\n",
        "\n",
        "Now we'll import the necessary libraries and set up our configuration for the MICRON model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "5a70a986",
      "metadata": {
        "id": "5a70a986"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import torch\n",
        "from pyhealth.datasets import MIMIC3Dataset\n",
        "from pyhealth.models import MICRON\n",
        "from pyhealth.trainer import Trainer\n",
        "# from pyhealth.metrics import multilabel_metrics # Removed this import\n",
        "\n",
        "# Set random seeds for reproducibility\n",
        "np.random.seed(42)\n",
        "torch.manual_seed(42)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed_all(42)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(f\"Using device: {DEVICE}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U5ALKIMXw6Ln",
        "outputId": "f3ce64f0-7370-4757-89c7-2da4fa2b48ad"
      },
      "id": "U5ALKIMXw6Ln",
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using device: cuda\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3577fe5a",
      "metadata": {
        "id": "3577fe5a"
      },
      "source": [
        "## 3. Load and Process MIMIC-III Dataset\n",
        "\n",
        "We'll load the MIMIC-III dataset using PyHealth's built-in dataset loader and prepare it for training. The dataset will be processed to include patient diagnoses, procedures, and medications."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dea4d6a5",
        "outputId": "31e4b3a6-0fbd-434a-a317-6586f019f523"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "No config path provided, using default config\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.mimic3:No config path provided, using default config\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Initializing mimic3 dataset from https://physionet.org/files/mimiciii-demo/1.4/ (dev mode: True)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Initializing mimic3 dataset from https://physionet.org/files/mimiciii-demo/1.4/ (dev mode: True)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Scanning table: patients from https://physionet.org/files/mimiciii-demo/1.4/PATIENTS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Scanning table: patients from https://physionet.org/files/mimiciii-demo/1.4/PATIENTS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/PATIENTS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/PATIENTS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Scanning table: admissions from https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Scanning table: admissions from https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Scanning table: icustays from https://physionet.org/files/mimiciii-demo/1.4/ICUSTAYS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Scanning table: icustays from https://physionet.org/files/mimiciii-demo/1.4/ICUSTAYS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ICUSTAYS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ICUSTAYS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Scanning table: diagnoses_icd from https://physionet.org/files/mimiciii-demo/1.4/DIAGNOSES_ICD.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Scanning table: diagnoses_icd from https://physionet.org/files/mimiciii-demo/1.4/DIAGNOSES_ICD.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/DIAGNOSES_ICD.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/DIAGNOSES_ICD.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Joining with table: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Joining with table: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Scanning table: procedures_icd from https://physionet.org/files/mimiciii-demo/1.4/PROCEDURES_ICD.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Scanning table: procedures_icd from https://physionet.org/files/mimiciii-demo/1.4/PROCEDURES_ICD.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/PROCEDURES_ICD.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/PROCEDURES_ICD.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Joining with table: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Joining with table: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Scanning table: prescriptions from https://physionet.org/files/mimiciii-demo/1.4/PRESCRIPTIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Scanning table: prescriptions from https://physionet.org/files/mimiciii-demo/1.4/PRESCRIPTIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/PRESCRIPTIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/PRESCRIPTIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Joining with table: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Joining with table: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Original path does not exist. Using alternative: https://physionet.org/files/mimiciii-demo/1.4/ADMISSIONS.csv\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting global event dataframe...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Collecting global event dataframe...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dev mode enabled: limiting to 1000 patients\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Dev mode enabled: limiting to 1000 patients\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collected dataframe with shape: (13030, 49)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Collected dataframe with shape: (13030, 49)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset: mimic3\n",
            "Dev mode: True\n",
            "Number of patients: 100\n",
            "Number of events: 13030\n"
          ]
        }
      ],
      "source": [
        "# Configuration\n",
        "#MIMIC3_PATH = \"https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III\"\n",
        "MIMIC3_PATH = \"https://physionet.org/files/mimiciii-demo/1.4/\"    #update this dataset path to your environment\n",
        "\n",
        "# Load MIMIC-III dataset\n",
        "dataset = MIMIC3Dataset(\n",
        "    root=MIMIC3_PATH,\n",
        "    tables=[\"DIAGNOSES_ICD\", \"PROCEDURES_ICD\", \"PRESCRIPTIONS\"],\n",
        "    dev=True\n",
        ")\n",
        "dataset.stats()"
      ],
      "id": "dea4d6a5"
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3. Set Drug Recommendation Task\n",
        "\n",
        "Use the DrugRecommendationMIMIC3 task function which creates samples with conditions, procedures, and atc-3 codes (drugs)."
      ],
      "metadata": {
        "id": "dbYRGd-ElhmE"
      },
      "id": "dbYRGd-ElhmE"
    },
    {
      "cell_type": "code",
      "source": [
        "from pyhealth.tasks import DrugRecommendationMIMIC3\n",
        "\n",
        "task = DrugRecommendationMIMIC3()\n",
        "samples = dataset.set_task(task, num_workers=4)\n",
        "\n",
        "print(f\"Sample Dataset Statistics:\")\n",
        "print(f\"\\t- Dataset: {samples.dataset_name}\")\n",
        "print(f\"\\t- Task: {samples.task_name}\")\n",
        "print(f\"\\t- Number of samples: {len(samples)}\")\n",
        "\n",
        "print(\"\\nFirst sample structure:\")\n",
        "print(f\"Patient ID: {samples.samples[0]['patient_id']}\")\n",
        "print(f\"Number of visits: {len(samples.samples[0]['conditions'])}\")\n",
        "print(f\"Sample conditions (first visit): {samples.samples[0]['conditions'][0][:5]}...\")\n",
        "print(f\"Sample procedures (first visit): {samples.samples[0]['procedures'][0][:5]}...\")\n",
        "print(f\"Sample drugs (target): {samples.samples[0]['drugs'][:10]}...\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HWnUXRYJlgK7",
        "outputId": "67c92f2c-976b-4286-faa0-bb5ac13ace5f"
      },
      "id": "HWnUXRYJlgK7",
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Setting task DrugRecommendationMIMIC3 for mimic3 base dataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Setting task DrugRecommendationMIMIC3 for mimic3 base dataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generating samples with 4 worker(s)...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Generating samples with 4 worker(s)...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generating samples for DrugRecommendationMIMIC3 with 4 workers\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.datasets.base_dataset:Generating samples for DrugRecommendationMIMIC3 with 4 workers\n",
            "Collecting samples for DrugRecommendationMIMIC3 from 4 workers: 100%|██████████| 100/100 [00:00<00:00, 1220.60it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Label drugs vocab: {'*NF*': 0, '0.45': 1, '0.9%': 2, '1/2 ': 3, '5% D': 4, 'AMP': 5, 'Acet': 6, 'Acyc': 7, 'Albu': 8, 'Alen': 9, 'Allo': 10, 'Alpr': 11, 'Alte': 12, 'Alum': 13, 'Amin': 14, 'Ampi': 15, 'Arti': 16, 'Asco': 17, 'Aspi': 18, 'Aten': 19, 'Ator': 20, 'Atov': 21, 'Azit': 22, 'Baci': 23, 'Bacl': 24, 'Bag': 25, 'Bisa': 26, 'Brim': 27, 'Bupi': 28, 'Cabe': 29, 'Calc': 30, 'Caph': 31, 'Caps': 32, 'Capt': 33, 'Carb': 34, 'CefT': 35, 'Cefa': 36, 'Cefe': 37, 'Ceft': 38, 'Cepa': 39, 'Chlo': 40, 'Cipr': 41, 'Cita': 42, 'Clin': 43, 'Clop': 44, 'Clot': 45, 'Coll': 46, 'Cosy': 47, 'Creo': 48, 'Crom': 49, 'Cycl': 50, 'Cypr': 51, 'Cyta': 52, 'D5 1': 53, 'D5NS': 54, 'D5W': 55, 'DOBU': 56, 'DOPa': 57, 'DOXO': 58, 'Daps': 59, 'Dapt': 60, 'Desi': 61, 'Dexa': 62, 'Dext': 63, 'Diaz': 64, 'Dilt': 65, 'Diph': 66, 'Docu': 67, 'Dola': 68, 'Done': 69, 'DopA': 70, 'Dorz': 71, 'Dost': 72, 'Doxy': 73, 'Dulo': 74, 'Enal': 75, 'Enox': 76, 'Epid': 77, 'Epoe': 78, 'Eryt': 79, 'Etop': 80, 'Famo': 81, 'Fat ': 82, 'Fent': 83, 'Ferr': 84, 'Filg': 85, 'Flee': 86, 'Fluc': 87, 'Flud': 88, 'Flut': 89, 'FoLI': 90, 'Fosp': 91, 'Furo': 92, 'Gaba': 93, 'Gamm': 94, 'Gelc': 95, 'Glip': 96, 'Gluc': 97, 'Guai': 98, 'HYDR': 99, 'Halo': 100, 'Hepa': 101, 'Hesp': 102, 'Humu': 103, 'Hydr': 104, 'Indo': 105, 'Infl': 106, 'Insu': 107, 'Ipra': 108, 'Iso-': 109, 'Isos': 110, 'Keto': 111, 'LR': 112, 'Lact': 113, 'Lans': 114, 'LeVE': 115, 'Leuc': 116, 'Levo': 117, 'Lido': 118, 'Line': 119, 'Lisi': 120, 'Lora': 121, 'Lumi': 122, 'Maal': 123, 'Magn': 124, 'Mege': 125, 'Mero': 126, 'Mesn': 127, 'MetR': 128, 'Meth': 129, 'Meto': 130, 'Metr': 131, 'Mico': 132, 'Mida': 133, 'Mido': 134, 'Milk': 135, 'Milr': 136, 'Mine': 137, 'Mirt': 138, 'Morp': 139, 'Mult': 140, 'NORe': 141, 'NS': 142, 'NS (': 143, 'Nado': 144, 'Nafc': 145, 'Nalo': 146, 'Naph': 147, 'Neom': 148, 'Neut': 149, 'Nitr': 150, 'Nore': 151, 'Nort': 152, 'Nyst': 153, 'Octr': 154, 'Olan': 155, 'Omep': 156, 'Onda': 157, 'Oxyc': 158, 'PHEN': 159, 'Panc': 160, 'Pant': 161, 'Paro': 162, 'Pent': 163, 'Phen': 164, 'Phos': 165, 'Phyt': 166, 'Pipe': 167, 'Pneu': 168, 'Poly': 169, 'Pota': 170, 'Pred': 171, 'Preg': 172, 'Pris': 173, 'Proc': 174, 'Prop': 175, 'Rani': 176, 'Rifa': 177, 'Ritu': 178, 'SW': 179, 'Sarn': 180, 'Scop': 181, 'Senn': 182, 'Sime': 183, 'Simv': 184, 'Siro': 185, 'Sodi': 186, 'Soln': 187, 'Spir': 188, 'Ster': 189, 'Sucr': 190, 'Sulf': 191, 'Syri': 192, 'Tacr': 193, 'Tams': 194, 'Thia': 195, 'Timo': 196, 'Tiot': 197, 'Tiza': 198, 'Tobr': 199, 'Tolt': 200, 'TraM': 201, 'Unas': 202, 'Vanc': 203, 'Vaso': 204, 'Vial': 205, 'Viga': 206, 'VinC': 207, 'Vita': 208, 'Warf': 209, 'Xope': 210, 'Zinc': 211, 'Zolp': 212, 'done': 213, 'sodi': 214, 'traZ': 215}\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "INFO:pyhealth.processors.label_processor:Label drugs vocab: {'*NF*': 0, '0.45': 1, '0.9%': 2, '1/2 ': 3, '5% D': 4, 'AMP': 5, 'Acet': 6, 'Acyc': 7, 'Albu': 8, 'Alen': 9, 'Allo': 10, 'Alpr': 11, 'Alte': 12, 'Alum': 13, 'Amin': 14, 'Ampi': 15, 'Arti': 16, 'Asco': 17, 'Aspi': 18, 'Aten': 19, 'Ator': 20, 'Atov': 21, 'Azit': 22, 'Baci': 23, 'Bacl': 24, 'Bag': 25, 'Bisa': 26, 'Brim': 27, 'Bupi': 28, 'Cabe': 29, 'Calc': 30, 'Caph': 31, 'Caps': 32, 'Capt': 33, 'Carb': 34, 'CefT': 35, 'Cefa': 36, 'Cefe': 37, 'Ceft': 38, 'Cepa': 39, 'Chlo': 40, 'Cipr': 41, 'Cita': 42, 'Clin': 43, 'Clop': 44, 'Clot': 45, 'Coll': 46, 'Cosy': 47, 'Creo': 48, 'Crom': 49, 'Cycl': 50, 'Cypr': 51, 'Cyta': 52, 'D5 1': 53, 'D5NS': 54, 'D5W': 55, 'DOBU': 56, 'DOPa': 57, 'DOXO': 58, 'Daps': 59, 'Dapt': 60, 'Desi': 61, 'Dexa': 62, 'Dext': 63, 'Diaz': 64, 'Dilt': 65, 'Diph': 66, 'Docu': 67, 'Dola': 68, 'Done': 69, 'DopA': 70, 'Dorz': 71, 'Dost': 72, 'Doxy': 73, 'Dulo': 74, 'Enal': 75, 'Enox': 76, 'Epid': 77, 'Epoe': 78, 'Eryt': 79, 'Etop': 80, 'Famo': 81, 'Fat ': 82, 'Fent': 83, 'Ferr': 84, 'Filg': 85, 'Flee': 86, 'Fluc': 87, 'Flud': 88, 'Flut': 89, 'FoLI': 90, 'Fosp': 91, 'Furo': 92, 'Gaba': 93, 'Gamm': 94, 'Gelc': 95, 'Glip': 96, 'Gluc': 97, 'Guai': 98, 'HYDR': 99, 'Halo': 100, 'Hepa': 101, 'Hesp': 102, 'Humu': 103, 'Hydr': 104, 'Indo': 105, 'Infl': 106, 'Insu': 107, 'Ipra': 108, 'Iso-': 109, 'Isos': 110, 'Keto': 111, 'LR': 112, 'Lact': 113, 'Lans': 114, 'LeVE': 115, 'Leuc': 116, 'Levo': 117, 'Lido': 118, 'Line': 119, 'Lisi': 120, 'Lora': 121, 'Lumi': 122, 'Maal': 123, 'Magn': 124, 'Mege': 125, 'Mero': 126, 'Mesn': 127, 'MetR': 128, 'Meth': 129, 'Meto': 130, 'Metr': 131, 'Mico': 132, 'Mida': 133, 'Mido': 134, 'Milk': 135, 'Milr': 136, 'Mine': 137, 'Mirt': 138, 'Morp': 139, 'Mult': 140, 'NORe': 141, 'NS': 142, 'NS (': 143, 'Nado': 144, 'Nafc': 145, 'Nalo': 146, 'Naph': 147, 'Neom': 148, 'Neut': 149, 'Nitr': 150, 'Nore': 151, 'Nort': 152, 'Nyst': 153, 'Octr': 154, 'Olan': 155, 'Omep': 156, 'Onda': 157, 'Oxyc': 158, 'PHEN': 159, 'Panc': 160, 'Pant': 161, 'Paro': 162, 'Pent': 163, 'Phen': 164, 'Phos': 165, 'Phyt': 166, 'Pipe': 167, 'Pneu': 168, 'Poly': 169, 'Pota': 170, 'Pred': 171, 'Preg': 172, 'Pris': 173, 'Proc': 174, 'Prop': 175, 'Rani': 176, 'Rifa': 177, 'Ritu': 178, 'SW': 179, 'Sarn': 180, 'Scop': 181, 'Senn': 182, 'Sime': 183, 'Simv': 184, 'Siro': 185, 'Sodi': 186, 'Soln': 187, 'Spir': 188, 'Ster': 189, 'Sucr': 190, 'Sulf': 191, 'Syri': 192, 'Tacr': 193, 'Tams': 194, 'Thia': 195, 'Timo': 196, 'Tiot': 197, 'Tiza': 198, 'Tobr': 199, 'Tolt': 200, 'TraM': 201, 'Unas': 202, 'Vanc': 203, 'Vaso': 204, 'Vial': 205, 'Viga': 206, 'VinC': 207, 'Vita': 208, 'Warf': 209, 'Xope': 210, 'Zinc': 211, 'Zolp': 212, 'done': 213, 'sodi': 214, 'traZ': 215}\n",
            "Processing samples: 100%|██████████| 36/36 [00:00<00:00, 1004.59it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated 36 samples for task DrugRecommendationMIMIC3\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "INFO:pyhealth.datasets.base_dataset:Generated 36 samples for task DrugRecommendationMIMIC3\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sample Dataset Statistics:\n",
            "\t- Dataset: mimic3\n",
            "\t- Task: <pyhealth.tasks.drug_recommendation.DrugRecommendationMIMIC3 object at 0x7f075ff85f70>\n",
            "\t- Number of samples: 36\n",
            "\n",
            "First sample structure:\n",
            "Patient ID: 40124\n",
            "Number of visits: 1\n",
            "Sample conditions (first visit): tensor([1, 2, 3, 4, 5])...\n",
            "Sample procedures (first visit): tensor([1, 0, 0, 0, 0])...\n",
            "Sample drugs (target): tensor([0., 0., 1., 0., 0., 0., 1., 0., 1., 0.])...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Split Dataset and Create Data Loaders"
      ],
      "metadata": {
        "id": "JW9OWmu9pfNp"
      },
      "id": "JW9OWmu9pfNp"
    },
    {
      "cell_type": "code",
      "source": [
        "from pyhealth.datasets import split_by_patient, get_dataloader\n",
        "\n",
        "train_dataset, val_dataset, test_dataset = split_by_patient(\n",
        "    samples, ratios=[0.7, 0.1, 0.2]\n",
        ")\n",
        "\n",
        "print(f\"Train samples: {len(train_dataset)}\")\n",
        "print(f\"Validation samples: {len(val_dataset)}\")\n",
        "print(f\"Test samples: {len(test_dataset)}\")\n",
        "\n",
        "train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)\n",
        "val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)\n",
        "test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xNq4dRmrtCWL",
        "outputId": "9aa706d9-c434-49fc-f881-9d0629242ef7"
      },
      "id": "xNq4dRmrtCWL",
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train samples: 27\n",
            "Validation samples: 3\n",
            "Test samples: 6\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b863949d",
      "metadata": {
        "id": "b863949d"
      },
      "source": [
        "## 4. Initialize and Configure MICRON Model\n",
        "\n",
        "Now we'll set up the MICRON model with appropriate hyperparameters for drug recommendation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "e0e2955c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e0e2955c",
        "outputId": "d8cfddf6-ff08-4b8e-e275-5b77c47aafcb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MICRON(\n",
            "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
            "    (conditions): Embedding(229, 128)\n",
            "    (procedures): Embedding(59, 128)\n",
            "    (drugs_hist): Embedding(198, 128)\n",
            "  ))\n",
            "  (micron): MICRONLayer(\n",
            "    (health_net): Linear(in_features=384, out_features=128, bias=True)\n",
            "    (prescription_net): Linear(in_features=128, out_features=128, bias=True)\n",
            "    (fc): Linear(in_features=128, out_features=216, bias=True)\n",
            "    (bce_loss_fn): BCEWithLogitsLoss()\n",
            "  )\n",
            ")\n",
            "MICRON(\n",
            "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
            "    (conditions): Embedding(229, 128)\n",
            "    (procedures): Embedding(59, 128)\n",
            "    (drugs_hist): Embedding(198, 128)\n",
            "  ))\n",
            "  (micron): MICRONLayer(\n",
            "    (health_net): Linear(in_features=384, out_features=128, bias=True)\n",
            "    (prescription_net): Linear(in_features=128, out_features=128, bias=True)\n",
            "    (fc): Linear(in_features=128, out_features=216, bias=True)\n",
            "    (bce_loss_fn): BCEWithLogitsLoss()\n",
            "  )\n",
            ")\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:MICRON(\n",
            "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
            "    (conditions): Embedding(229, 128)\n",
            "    (procedures): Embedding(59, 128)\n",
            "    (drugs_hist): Embedding(198, 128)\n",
            "  ))\n",
            "  (micron): MICRONLayer(\n",
            "    (health_net): Linear(in_features=384, out_features=128, bias=True)\n",
            "    (prescription_net): Linear(in_features=128, out_features=128, bias=True)\n",
            "    (fc): Linear(in_features=128, out_features=216, bias=True)\n",
            "    (bce_loss_fn): BCEWithLogitsLoss()\n",
            "  )\n",
            ")\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Metrics: ['jaccard_samples', 'f1_samples', 'pr_auc_samples', 'ddi']\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Metrics: ['jaccard_samples', 'f1_samples', 'pr_auc_samples', 'ddi']\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Device: cuda\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Device: cuda\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Baseline performance before training:\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 204.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'jaccard_samples': 0.14468260340523165, 'f1_samples': 0.25038132615979186, 'pr_auc_samples': 0.16410485822712173, 'ddi_score': 0.0, 'loss': 26.26311683654785}\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# Model hyperparameters\n",
        "model_params = {\n",
        "    \"embedding_dim\": 128,\n",
        "    \"hidden_dim\": 128,\n",
        "    \"lam\": 0.1  # Regularization parameter for reconstruction loss\n",
        "}\n",
        "\n",
        "# Initialize MICRON model\n",
        "model = MICRON(\n",
        "    dataset=samples,\n",
        "    **model_params\n",
        ").to(DEVICE)\n",
        "\n",
        "print(model)\n",
        "\n",
        "\n",
        "# Configure trainer\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    #metrics=[\"jaccard_samples\", \"f1_samples\", \"pr_auc_samples\"]\n",
        "    metrics=[\"jaccard_samples\", \"f1_samples\", \"pr_auc_samples\", \"ddi\"]\n",
        ")\n",
        "\n",
        "print(\"Baseline performance before training:\")\n",
        "baseline_results = trainer.evaluate(test_dataloader)\n",
        "print(baseline_results)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "XcZI0ohus59W"
      },
      "id": "XcZI0ohus59W"
    },
    {
      "cell_type": "markdown",
      "id": "209841fe",
      "metadata": {
        "id": "209841fe"
      },
      "source": [
        "## 5. Train the Model\n",
        "\n",
        "Let's train the MICRON model on our processed MIMIC-III dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "d2b6610d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "56243ebc60f347cf94142d9eb954b92c",
            "c8e14ca23e9d45bcbd8abfc3804b4912",
            "8ca24a121e6344c194115c6b5d6a11b7",
            "40c6a345bc2044aab5d301baa501df14",
            "a6ae311f1aa24dc5805026971fe56d82",
            "99b03d2966d143efbac576bd3b97c42b",
            "87680847c231446c9f975fbd367c54b6",
            "5ac37ba2c2724aa782ae7b3211861520",
            "3a50a1d9f1224ba0b747bba5bef9a9f9",
            "dab14088fc3b419fbfa966bfbd8ed0b5",
            "39df669e3ce446d6b54f1718ce149398",
            "dd0c9925874a411fb9ff19a98f0d59ab",
            "1053130b25054630822d3b4f29a66592",
            "27595c61f7e5466b841bfccb9858aed0",
            "c58ed5800ff54798963ff06b014a3b1f",
            "0fb6f47c55d94d438ea8c627ba36ba59",
            "d997c3ce4fcc46a1aed803db4f13ad75",
            "fb02bb52b14d43c18b3e06881cea3fc6",
            "4beabf36b97a47c8bb29a64cb86fe05d",
            "f46f5f9891564dffac6581520dad0e58",
            "b8bed2150ed44200930e808d739da8e3",
            "1ec681d0c5c741cd9f80b0fd3f00326c",
            "dd3227a1a73941a2ae121bb4fc6d3a2b",
            "609d1928ac2a4d1e9e8580a11872396f",
            "306f260c1c75495a895adeb1c559359f",
            "170060f6929b4c308f3c57682d5f4cb6",
            "5f43853543844784926b45f3e02057c9",
            "80c084adcf9a489591813440dee42021",
            "86cc5e05adb24392b6c36e1156874d4f",
            "b696a0c6bdc7456399806023b36a9d3b",
            "aeb41ccfb61e4dfdaefc7d1dfb3aedbf",
            "2396f96cdfc4488694b2e5ad0ddb7e61",
            "ded06b52a2b94ef89ae9bfa594d7bec9",
            "f263ff1135d54563bd2da89fe8de7fe7",
            "3b912281619743a0b3f8c7527210cd60",
            "4799cd0fb2cb47aca1d5b7dc531e7e08",
            "2cf42af9f2684fb89c5cac306491c459",
            "701ebc13a90f44b196817d0d894aa554",
            "ac08faae67f74051a52a7d9e179131df",
            "0a95232677024ea59c8560323ab54437",
            "863b09f4558e444e9d8f170d5476d223",
            "31f4f48e77464c28bada16c192b5a0d4",
            "b1314a7dcad344dd8ed63402fe297312",
            "3e312fca700c49fd86c11af4436014a8",
            "c9a00f0dc63c42f59da773d570cf0dbb",
            "6986dc7c27374b5f8064deb826c41c28",
            "40bc0757512c4d1a8096790588d966f5",
            "9ca609925764491d94269ee6d41ec161",
            "4c2b818383e9431cb8dcbdce1d8eca8d",
            "0f2346468406416d9ea29b66267d443b",
            "ceebf6655f46451384896aed98370504",
            "753e1e3515014e36a1730b19e08c58d2",
            "27475eddc9c2469b95a5458bebc422b2",
            "73366005b06c4e2299c219ac29723a39",
            "120721fff18c4afb8d94bee41e73e191"
          ]
        },
        "id": "d2b6610d",
        "outputId": "385ca2e8-745c-482b-a4a0-c7877c88ad59"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training:\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Training:\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch size: 32\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Batch size: 32\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizer: <class 'torch.optim.adam.Adam'>\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Optimizer: <class 'torch.optim.adam.Adam'>\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizer params: {'lr': 0.001}\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Optimizer params: {'lr': 0.001}\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Weight decay: 0.0\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Weight decay: 0.0\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Max grad norm: None\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Max grad norm: None\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x7f05e45c9d00>\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x7f05e45c9d00>\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Monitor: pr_auc_samples\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Monitor: pr_auc_samples\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Monitor criterion: max\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Monitor criterion: max\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epochs: 5\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Epochs: 5\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Epoch 0 / 5:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "56243ebc60f347cf94142d9eb954b92c"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Train epoch-0, step-1 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Train epoch-0, step-1 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 26.7790\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 26.7790\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 261.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Eval epoch-0, step-1 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "INFO:pyhealth.trainer:--- Eval epoch-0, step-1 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jaccard_samples: 0.1060\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:jaccard_samples: 0.1060\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "f1_samples: 0.1902\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:f1_samples: 0.1902\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "pr_auc_samples: 0.1278\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:pr_auc_samples: 0.1278\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 12.1593\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 12.1593\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "New best pr_auc_samples score (0.1278) at epoch-0, step-1\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:New best pr_auc_samples score (0.1278) at epoch-0, step-1\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Epoch 1 / 5:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "dd0c9925874a411fb9ff19a98f0d59ab"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Train epoch-1, step-2 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Train epoch-1, step-2 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 12.1609\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 12.1609\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 231.73it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Eval epoch-1, step-2 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Eval epoch-1, step-2 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jaccard_samples: 0.1043\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:jaccard_samples: 0.1043\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "f1_samples: 0.1881\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:f1_samples: 0.1881\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "pr_auc_samples: 0.1253\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:pr_auc_samples: 0.1253\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 8.5477\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 8.5477\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Epoch 2 / 5:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "dd3227a1a73941a2ae121bb4fc6d3a2b"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Train epoch-2, step-3 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Train epoch-2, step-3 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 8.3976\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 8.3976\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 235.13it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Eval epoch-2, step-3 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Eval epoch-2, step-3 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jaccard_samples: 0.0841\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:jaccard_samples: 0.0841\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "f1_samples: 0.1539\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:f1_samples: 0.1539\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "pr_auc_samples: 0.1289\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:pr_auc_samples: 0.1289\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 8.2150\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 8.2150\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "New best pr_auc_samples score (0.1289) at epoch-2, step-3\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:New best pr_auc_samples score (0.1289) at epoch-2, step-3\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Epoch 3 / 5:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f263ff1135d54563bd2da89fe8de7fe7"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Train epoch-3, step-4 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Train epoch-3, step-4 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 8.2572\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 8.2572\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 225.23it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Eval epoch-3, step-4 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Eval epoch-3, step-4 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jaccard_samples: 0.1594\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:jaccard_samples: 0.1594\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "f1_samples: 0.2718\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:f1_samples: 0.2718\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "pr_auc_samples: 0.1652\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:pr_auc_samples: 0.1652\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 6.2571\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 6.2571\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "New best pr_auc_samples score (0.1652) at epoch-3, step-4\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:New best pr_auc_samples score (0.1652) at epoch-3, step-4\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Epoch 4 / 5:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "c9a00f0dc63c42f59da773d570cf0dbb"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Train epoch-4, step-5 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Train epoch-4, step-5 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 6.6966\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 6.6966\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 266.20it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--- Eval epoch-4, step-5 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:--- Eval epoch-4, step-5 ---\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jaccard_samples: 0.1548\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:jaccard_samples: 0.1548\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "f1_samples: 0.2655\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:f1_samples: 0.2655\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "pr_auc_samples: 0.1809\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:pr_auc_samples: 0.1809\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:ddi_score: 0.0000\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss: 5.0584\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:loss: 5.0584\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "New best pr_auc_samples score (0.1809) at epoch-4, step-5\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:New best pr_auc_samples score (0.1809) at epoch-4, step-5\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loaded best model\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pyhealth.trainer:Loaded best model\n"
          ]
        }
      ],
      "source": [
        "# Train the model\n",
        "history = trainer.train(\n",
        "    train_dataloader, # Pass as positional argument\n",
        "    val_dataloader,   # Pass as positional argument\n",
        "    epochs=5,\n",
        "    monitor=\"pr_auc_samples\"\n",
        ")\n",
        "\n",
        "# Save the trained model\n",
        "#torch.save(model.state_dict(), \"/content/drive/MyDrive/micron_model.pt\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "384e52d9",
      "metadata": {
        "id": "384e52d9"
      },
      "source": [
        "## 6. Evaluate Model Performance\n",
        "\n",
        "Finally, let's evaluate our trained model on the test set and visualize the results."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test_results = trainer.evaluate(test_dataloader)\n",
        "print(\"Final test set performance:\")\n",
        "print(test_results)\n",
        "\n",
        "print(f\"\\nKey Metrics:\")\n",
        "print(f\"  PR-AUC: {test_results.get('pr_auc_samples', 'N/A'):.4f}\")\n",
        "print(f\"  F1 Score: {test_results.get('f1_samples', 'N/A'):.4f}\")\n",
        "print(f\"  Jaccard: {test_results.get('jaccard_samples', 'N/A'):.4f}\")\n",
        "print(f\"  DDI Rate: {test_results.get('ddi_score', 'N/A'):.4f} (lower is better)\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "muPnVGey13pl",
        "outputId": "9905a1a4-0e01-4147-86f2-9df6fea4d8a6"
      },
      "id": "muPnVGey13pl",
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 197.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Final test set performance:\n",
            "{'jaccard_samples': 0.20728042582268005, 'f1_samples': 0.33950176113659936, 'pr_auc_samples': 0.27233906976425537, 'ddi_score': 0.0, 'loss': 4.690524578094482}\n",
            "\n",
            "Key Metrics:\n",
            "  PR-AUC: 0.2723\n",
            "  F1 Score: 0.3395\n",
            "  Jaccard: 0.2073\n",
            "  DDI Rate: 0.0000 (lower is better)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# before training\n",
        "# {'jaccard_samples': 0.14468260340523165, 'f1_samples': 0.25038132615979186, 'pr_auc_samples': 0.16410485822712173, 'ddi_score': 0.0, 'loss': 26.26311683654785}\n",
        "\n",
        "# Post training\n",
        "# {'jaccard_samples': 0.20728042582268005, 'f1_samples': 0.33950176113659936, 'pr_auc_samples': 0.27233906976425537, 'ddi_score': 0.0, 'loss': 4.690524578094482}"
      ],
      "metadata": {
        "id": "QeZ_64qL2SnI"
      },
      "id": "QeZ_64qL2SnI",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "47925060",
      "metadata": {
        "id": "47925060"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "We have successfully implemented and trained a MICRON model for drug recommendation using the MIMIC-III dataset. The model's performance can be evaluated using the metrics above:\n",
        "\n",
        "1. DDI score: Shows the drug-drug interaction score indicating the effect of one drug on another. lower is better.\n",
        "2. Precision: Indicates how many of the predicted drugs were actually correct\n",
        "3. Recall: Shows how many of the actual drugs were correctly predicted\n",
        "4. F1 Score: The harmonic mean of precision and recall\n",
        "\n",
        "The confusion matrix visualization helps us understand where the model performs well and where it might need improvement. The training loss plot shows how the model learned over time.\n",
        "\n",
        "Next steps could include:\n",
        "- Hyperparameter tuning to improve performance\n",
        "- Testing with different model architectures\n",
        "- Analyzing specific cases where the model performs well or poorly\n",
        "- Incorporating additional patient features"
      ]
    }
  ],
  "metadata": {
    "language_info": {
      "name": "python"
    },
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "56243ebc60f347cf94142d9eb954b92c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_c8e14ca23e9d45bcbd8abfc3804b4912",
              "IPY_MODEL_8ca24a121e6344c194115c6b5d6a11b7",
              "IPY_MODEL_40c6a345bc2044aab5d301baa501df14"
            ],
            "layout": "IPY_MODEL_a6ae311f1aa24dc5805026971fe56d82"
          }
        },
        "c8e14ca23e9d45bcbd8abfc3804b4912": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_99b03d2966d143efbac576bd3b97c42b",
            "placeholder": "​",
            "style": "IPY_MODEL_87680847c231446c9f975fbd367c54b6",
            "value": "Epoch 0 / 5: 100%"
          }
        },
        "8ca24a121e6344c194115c6b5d6a11b7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5ac37ba2c2724aa782ae7b3211861520",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_3a50a1d9f1224ba0b747bba5bef9a9f9",
            "value": 1
          }
        },
        "40c6a345bc2044aab5d301baa501df14": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_dab14088fc3b419fbfa966bfbd8ed0b5",
            "placeholder": "​",
            "style": "IPY_MODEL_39df669e3ce446d6b54f1718ce149398",
            "value": " 1/1 [00:00&lt;00:00,  4.81it/s]"
          }
        },
        "a6ae311f1aa24dc5805026971fe56d82": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "99b03d2966d143efbac576bd3b97c42b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "87680847c231446c9f975fbd367c54b6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "5ac37ba2c2724aa782ae7b3211861520": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3a50a1d9f1224ba0b747bba5bef9a9f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "dab14088fc3b419fbfa966bfbd8ed0b5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "39df669e3ce446d6b54f1718ce149398": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "dd0c9925874a411fb9ff19a98f0d59ab": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_1053130b25054630822d3b4f29a66592",
              "IPY_MODEL_27595c61f7e5466b841bfccb9858aed0",
              "IPY_MODEL_c58ed5800ff54798963ff06b014a3b1f"
            ],
            "layout": "IPY_MODEL_0fb6f47c55d94d438ea8c627ba36ba59"
          }
        },
        "1053130b25054630822d3b4f29a66592": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d997c3ce4fcc46a1aed803db4f13ad75",
            "placeholder": "​",
            "style": "IPY_MODEL_fb02bb52b14d43c18b3e06881cea3fc6",
            "value": "Epoch 1 / 5: 100%"
          }
        },
        "27595c61f7e5466b841bfccb9858aed0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4beabf36b97a47c8bb29a64cb86fe05d",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_f46f5f9891564dffac6581520dad0e58",
            "value": 1
          }
        },
        "c58ed5800ff54798963ff06b014a3b1f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b8bed2150ed44200930e808d739da8e3",
            "placeholder": "​",
            "style": "IPY_MODEL_1ec681d0c5c741cd9f80b0fd3f00326c",
            "value": " 1/1 [00:00&lt;00:00, 23.98it/s]"
          }
        },
        "0fb6f47c55d94d438ea8c627ba36ba59": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d997c3ce4fcc46a1aed803db4f13ad75": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fb02bb52b14d43c18b3e06881cea3fc6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4beabf36b97a47c8bb29a64cb86fe05d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f46f5f9891564dffac6581520dad0e58": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "b8bed2150ed44200930e808d739da8e3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1ec681d0c5c741cd9f80b0fd3f00326c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "dd3227a1a73941a2ae121bb4fc6d3a2b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_609d1928ac2a4d1e9e8580a11872396f",
              "IPY_MODEL_306f260c1c75495a895adeb1c559359f",
              "IPY_MODEL_170060f6929b4c308f3c57682d5f4cb6"
            ],
            "layout": "IPY_MODEL_5f43853543844784926b45f3e02057c9"
          }
        },
        "609d1928ac2a4d1e9e8580a11872396f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_80c084adcf9a489591813440dee42021",
            "placeholder": "​",
            "style": "IPY_MODEL_86cc5e05adb24392b6c36e1156874d4f",
            "value": "Epoch 2 / 5: 100%"
          }
        },
        "306f260c1c75495a895adeb1c559359f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b696a0c6bdc7456399806023b36a9d3b",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_aeb41ccfb61e4dfdaefc7d1dfb3aedbf",
            "value": 1
          }
        },
        "170060f6929b4c308f3c57682d5f4cb6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_2396f96cdfc4488694b2e5ad0ddb7e61",
            "placeholder": "​",
            "style": "IPY_MODEL_ded06b52a2b94ef89ae9bfa594d7bec9",
            "value": " 1/1 [00:00&lt;00:00, 26.16it/s]"
          }
        },
        "5f43853543844784926b45f3e02057c9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "80c084adcf9a489591813440dee42021": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "86cc5e05adb24392b6c36e1156874d4f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b696a0c6bdc7456399806023b36a9d3b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "aeb41ccfb61e4dfdaefc7d1dfb3aedbf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "2396f96cdfc4488694b2e5ad0ddb7e61": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ded06b52a2b94ef89ae9bfa594d7bec9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "f263ff1135d54563bd2da89fe8de7fe7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_3b912281619743a0b3f8c7527210cd60",
              "IPY_MODEL_4799cd0fb2cb47aca1d5b7dc531e7e08",
              "IPY_MODEL_2cf42af9f2684fb89c5cac306491c459"
            ],
            "layout": "IPY_MODEL_701ebc13a90f44b196817d0d894aa554"
          }
        },
        "3b912281619743a0b3f8c7527210cd60": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ac08faae67f74051a52a7d9e179131df",
            "placeholder": "​",
            "style": "IPY_MODEL_0a95232677024ea59c8560323ab54437",
            "value": "Epoch 3 / 5: 100%"
          }
        },
        "4799cd0fb2cb47aca1d5b7dc531e7e08": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_863b09f4558e444e9d8f170d5476d223",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_31f4f48e77464c28bada16c192b5a0d4",
            "value": 1
          }
        },
        "2cf42af9f2684fb89c5cac306491c459": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b1314a7dcad344dd8ed63402fe297312",
            "placeholder": "​",
            "style": "IPY_MODEL_3e312fca700c49fd86c11af4436014a8",
            "value": " 1/1 [00:00&lt;00:00, 26.06it/s]"
          }
        },
        "701ebc13a90f44b196817d0d894aa554": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ac08faae67f74051a52a7d9e179131df": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0a95232677024ea59c8560323ab54437": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "863b09f4558e444e9d8f170d5476d223": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "31f4f48e77464c28bada16c192b5a0d4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "b1314a7dcad344dd8ed63402fe297312": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3e312fca700c49fd86c11af4436014a8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "c9a00f0dc63c42f59da773d570cf0dbb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_6986dc7c27374b5f8064deb826c41c28",
              "IPY_MODEL_40bc0757512c4d1a8096790588d966f5",
              "IPY_MODEL_9ca609925764491d94269ee6d41ec161"
            ],
            "layout": "IPY_MODEL_4c2b818383e9431cb8dcbdce1d8eca8d"
          }
        },
        "6986dc7c27374b5f8064deb826c41c28": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0f2346468406416d9ea29b66267d443b",
            "placeholder": "​",
            "style": "IPY_MODEL_ceebf6655f46451384896aed98370504",
            "value": "Epoch 4 / 5: 100%"
          }
        },
        "40bc0757512c4d1a8096790588d966f5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_753e1e3515014e36a1730b19e08c58d2",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_27475eddc9c2469b95a5458bebc422b2",
            "value": 1
          }
        },
        "9ca609925764491d94269ee6d41ec161": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_73366005b06c4e2299c219ac29723a39",
            "placeholder": "​",
            "style": "IPY_MODEL_120721fff18c4afb8d94bee41e73e191",
            "value": " 1/1 [00:00&lt;00:00, 26.30it/s]"
          }
        },
        "4c2b818383e9431cb8dcbdce1d8eca8d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0f2346468406416d9ea29b66267d443b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ceebf6655f46451384896aed98370504": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "753e1e3515014e36a1730b19e08c58d2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "27475eddc9c2469b95a5458bebc422b2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "73366005b06c4e2299c219ac29723a39": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "120721fff18c4afb8d94bee41e73e191": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}